package org.dreamwork.tools.sherpa.ncnn.examples;

import com.k2fsa.sherpa.ncnn.*;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

public class SherpaHelper {
    public static SherpaNcnn initRecognizer (int type) {
        System.out.print ("initialing recognizer ... ");
        long start = System.currentTimeMillis ();
        FeatureExtractorConfig featConfig = new FeatureExtractorConfig (16000.0f, 80);
        ModelConfig modelConfig = ModelConfig.getModelConfig (type, true);
        DecoderConfig decoderConfig = new DecoderConfig ("greedy_search", 4);
        RecognizerConfig recognizerConfig = new RecognizerConfig (
                featConfig, modelConfig, decoderConfig,
                true, 2.4f, 1.4f, 20f,
                "", 1.5f
        );
        SherpaNcnn sherpa = new SherpaNcnn (recognizerConfig);
        long end = System.currentTimeMillis ();
        System.out.printf ("done, it takes %.3f seconds.%n", (end - start) / 1000f);
        return sherpa;
    }

    public static void initNcnnModel (String... args) throws IOException {
        String modelDir = null;
        if (args.length > 0) {
            for (int i = 0; i < args.length; i ++) {
                String option = args[i];
                if ("-m".equals (option) && args.length > i + 1) {
                    modelDir = args [i + 1];
                    break;
                } else if (option.startsWith ("--model-dir=")) {
                    modelDir = option.substring ("--model-dir=".length ());
                    break;
                }
            }
        }
        if (modelDir == null || modelDir.isEmpty ()) {
            System.err.println ("model dir does not set, using the default location: " + new File (".").getAbsolutePath ());
        } else {
            Path path = Paths.get (modelDir);
            if (Files.notExists (path)) {
                System.err.println ("location " + path + " does not exists.");
                System.exit (-1);
            } else if (!Files.isDirectory (path)) {
                System.err.println ("location " + path.toRealPath () + " is not a directory.");
                System.exit (-1);
            } else if (!Files.isReadable (path)) {
                System.err.println ("directory " + path.toRealPath () + " is not readable.");
                System.exit (-1);
            } else {
                String realPath = path.toRealPath ().toString ();
                System.out.println ("using model directory: " + realPath);
                System.setProperty (ModelConfig.MODEL_DIR_KEY, realPath);
            }
        }
    }
}
